# your_analysis_code.py - 专为阿里云函数计算优化
import cv2
import numpy as np
import os
import json
import uuid
from ultralytics import YOLO
import mypose
import pandas as pd
from oss2 import Auth, Bucket
import logging

logger = logging.getLogger()

class ReportGenerator:
    def __init__(self, fps):
        # 创建一个列表来存储分析结果
        self.analysis_results = []
        self.excellent_count = 0  # 计数优秀的相似度
        self.fps = fps
        self.frame_count = 0  # 初始化帧计数
        self.excellent_count = 0

    def append_result(self, similarities, avg_similarity):
        # 记录当前时刻（以秒为单位）
        current_time = self.frame_count / self.fps
        self.frame_count += 1
        result = {
            "时刻": current_time,
            "整体平均相似度": avg_similarity,
        }
        # 添加每个关节的相似度
        for idx, angle in enumerate(similarities):
            result[f"关节{idx + 1}相似度"] = angle  # 使用索引作为关节编号
        self.analysis_results.append(result)  # 将结果添加到列表中
        # 统计优秀的相似度
        if avg_similarity > 0.9:  # 超过90%
            self.excellent_count += 1

    def generate_report(self, excel_path, json_path):
        """生成Excel报告和JSON数据文件"""
        if self.frame_count > 0:
            excellent_ratio = (
                self.excellent_count / self.frame_count
            ) * 100  # 计算优秀比例
            logger.info(f"优秀比例: {excellent_ratio:.2f}%")  # 记录日志
            
            # 创建DataFrame
            df = pd.DataFrame(self.analysis_results)
            
            # 保存为Excel
            df.to_excel(excel_path, index=False)
            
            # 保存为JSON（更易于前端读取）
            with open(json_path, 'w') as f:
                json.dump(self.analysis_results, f)
            
            return {
                "excel_path": excel_path,
                "json_path": json_path,
                "excellent_ratio": excellent_ratio,
                "total_frames": self.frame_count
            }
        return None

class DanceEvaluation:
    def __init__(self, video_path, output_path):
        logger.info(f"初始化舞蹈评估: {video_path} -> {output_path}")
        self.model = YOLO("yolov8n-pose.pt")  # 确保模型文件在部署包中
        self.cap = cv2.VideoCapture(video_path)
        
        if not self.cap.isOpened():
            raise ValueError(f"无法打开视频文件: {video_path}")
            
        self.frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.fps = self.cap.get(cv2.CAP_PROP_FPS)
        logger.info(f"视频参数: {self.frame_width}x{self.frame_height} @ {self.fps}fps")
        
        self.out_path = output_path
        self.fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        self.out = cv2.VideoWriter(
            self.out_path,
            self.fourcc,
            self.fps,
            (self.frame_width, self.frame_height),
        )
        self.report = ReportGenerator(self.fps)

    def evaluate_dance_frame(self, frame):
        """分析单帧并返回相似度"""
        results = self.model(frame)
        
        if len(results) >= 2:
            pose1 = mypose.PoseKeypoints(
                results[0], self.frame_width, self.frame_height, person_id=0
            )
            pose2 = mypose.PoseKeypoints(
                results[0], self.frame_width, self.frame_height, person_id=1
            )

            # 确定左右位置（基于x坐标总和）
            sumx_1 = sum([pose1.get_keypoint(idx)[0] for idx in range(16)])
            sumx_2 = sum([pose2.get_keypoint(idx)[0] for idx in range(16)])
            if sumx_1 > sumx_2:
                pose1, pose2 = pose2, pose1

            # 显示姿态
            pose1.PoseShow(frame, show_angles=True)
            pose2.PoseShow(frame, show_angles=True)

            # 比较姿态
            avg_similarity, similarities = mypose.compare_poses(frame, pose1, pose2)
            self.report.append_result(similarities, avg_similarity)
            return avg_similarity
        return None

    def evaluate_dance(self):
        """分析整个视频"""
        frame_count = 0
        total_similarity = 0
        valid_frames = 0
        
        logger.info("开始视频分析...")
        while True:
            ret, frame = self.cap.read()
            if not ret:
                break
                
            frame_count += 1
            if frame_count % 30 == 0:  # 每30帧记录一次
                logger.info(f"已处理 {frame_count} 帧")
                
            similarity = self.evaluate_dance_frame(frame)
            if similarity is not None:
                total_similarity += similarity
                valid_frames += 1
                
            self.out.write(frame)
        
        self.cap.release()
        self.out.release()
        
        avg_overall = total_similarity / valid_frames if valid_frames > 0 else 0
        logger.info(f"分析完成! 共处理 {frame_count} 帧, 有效帧 {valid_frames}, 平均相似度: {avg_overall:.2f}")
        
        return {
            "total_frames": frame_count,
            "valid_frames": valid_frames,
            "avg_similarity": avg_overall
        }

    def generate_report(self, excel_path, json_path):
        """生成报告"""
        return self.report.generate_report(excel_path, json_path)


def analyze_joined_video(video_path):
    """
    分析已拼接的双视频（左右分屏）
    返回包含标记视频和相似度数据的字典
    """
    logger.info(f"开始分析视频: {video_path}")
    
    # 1. 生成唯一任务ID
    task_id = f"task_{uuid.uuid4().hex[:8]}"
    
    # 2. 设置临时路径
    marked_video_path = f"/tmp/{task_id}_marked.mp4"
    excel_report_path = f"/tmp/{task_id}_report.xlsx"
    json_report_path = f"/tmp/{task_id}_report.json"
    
    try:
        # 3. 执行分析
        evaluation = DanceEvaluation(video_path, marked_video_path)
        stats = evaluation.evaluate_dance()
        report_data = evaluation.generate_report(excel_report_path, json_report_path)
        
        # 4. 准备返回结果
        result = {
            "task_id": task_id,
            "marked_video": marked_video_path,
            "excel_report": excel_report_path,
            "json_report": json_report_path,
            "stats": stats,
            "report_data": report_data
        }
        
        logger.info(f"分析成功完成: {task_id}")
        return result
        
    except Exception as e:
        logger.error(f"分析失败: {str(e)}", exc_info=True)
        raise e
